from PIL import Image
import numpy as np
import torchvision


class DataAugment(object):
    # 改需求时需要
    def __init__(self):
        super(DataAugment, self).__init__()

    # 确定size传入的是元组格式 (h,w)
    def detect_resize(self, img, label, size: tuple):
        # trans 是三个类方法实例
        trans = torchvision.transforms.Compose([
            # cvimg to PILimg
            torchvision.transforms.ToPILImage(),
            torchvision.transforms.Resize(size=size),
            torchvision.transforms.ToTensor()
        ])
        image = trans(img)
        # 目标检测时也需要进行label trans
        return image, label

    # 实例化
    # A call方法 实例化对象a使用a()来调用
    # A()(img) 不用实例化 eg torchvision.transforms.ToPILImage()(img)
    # args:tuple  kwargs:dict
    def __call__(self, *args, **kwargs):
        return self.detect_resize(*args)